package org.apache.lorry.zookeeper.cluster;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.lorry.zookeeper.meeting.ZkConstant;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.List;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

@Slf4j
public class Master {
    ZooKeeper zk;

    private volatile boolean connected = false;
    private volatile boolean expired = false;

    private final String ipPort;

    private String name;

    private String leaderId;

    private Status status = Status.RUNNING;


    private ChildCache workerCache;


    private ChildCache taskCache;

    private Random random = new Random(this.hashCode());

    static enum Status {
        RUNNING, ELECTED, NOTELECTED,
    }

    public static final int SESSION_TIMEOUT = 10000;

    public Master(String ipPort) {
        this.ipPort = ipPort;
    }

    public void startZK() throws Exception {
        CountDownLatch latch = new CountDownLatch(1);
        zk = new ZooKeeper(ipPort, SESSION_TIMEOUT, (event) -> {
            switch (event.getState()) {
                case SyncConnected:
                    connected = true;
                    latch.countDown();
                    break;
                case Disconnected:
                    connected = false;
                    log.error("Disconnected");
                    break;
                case Expired:
                    connected = false;
                    expired = true;
                    System.out.println("Expired");
                    break;
                default:
                    break;
            }
        });
        latch.await(SESSION_TIMEOUT, TimeUnit.MICROSECONDS);
    }


    public void init() throws InterruptedException, KeeperException {
        initCluster();
        this.name = genName();
    }

    public String name() {
        return name;
    }

    private String genName() throws InterruptedException, KeeperException {
        String namePath = zk.create(ZooConstant.SEQ_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT_SEQUENTIAL);
        return "NO_" + namePath.substring(ZooConstant.SEQ_PATH.length() + 1);
    }

    private void initCluster() throws InterruptedException, KeeperException {
        Stat stat = zk.exists(ZooConstant.ROOT_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.ROOT_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        stat = zk.exists(ZooConstant.WORKER_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.WORKER_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        stat = zk.exists(ZooConstant.TASK_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.TASK_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        stat = zk.exists(ZooConstant.ASSIGN_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.ASSIGN_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        stat = zk.exists(ZooConstant.STATUS_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.STATUS_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
    }

    public void runForMaster() {
        zk.create(ZooConstant.MASTER_PATH, name.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL, masterCreateCallback, name);
    }

    private AsyncCallback.StringCallback masterCreateCallback = (rc, path, ctx, name) -> {
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                checkMaster();  //如果连接丢失，检查Master是否存在
                return;
            case OK:
                log.info("I am master , the name is {}", this.name);
                this.leaderId = (String) ctx;
                status = Status.ELECTED;
                takeLeaderShip();
                break;
            case NODEEXISTS:
                log.info("master already exists");
                status = Status.NOTELECTED;
                masterExists();
                break;
            default:
                log.error("something went wrong", KeeperException.create(KeeperException.Code.get(rc), path));
        }
    };

    private void takeLeaderShip() {
        log.info("Going for list of workers");
        getWorkers();

        new RecoveredAssignments(zk).recover((rc, tasks) -> {
            if (rc == RecoveredAssignments.RecoveryCallback.FAILED) {
                log.error("Recovery of assigned tasks failed.");
            } else {
                log.info("Assigning recovered tasks");
                getTasks();
            }
        });
    }


    private void getTasks() {
        zk.getChildren(ZooConstant.TASK_PATH, tasksChangeWatcher, tasksGetChildrenCallback, null);
    }

    Watcher tasksChangeWatcher = (event) -> {
        if (event.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
            assert ZooConstant.TASK_PATH.equals(event.getPath());
            getTasks();
        }
    };

    AsyncCallback.ChildrenCallback tasksGetChildrenCallback = (rc, path, ctx, children) -> {
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                getTasks();
                break;
            case OK:
                log.info("Succesfully got a list of tasks: {}", children);
                List<String> toProcess;
                if (taskCache == null) {
                    taskCache = new ChildCache(children);
                    toProcess = children;
                }else{
                    toProcess = taskCache.addChildDiff(children);
                }

                if (toProcess != null) {
                    assignTasks(toProcess);
                }
                break;
            default:
                log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
        }
    };

    private void assignTasks(List<String> toProcess) {
        for (String task : toProcess) {
            getTaskData(task);
        }
    }

    private void getTaskData(String task) {
        zk.getData(ZooConstant.TASK_PATH + "/" + task, false, taskDataCallback, task);
    }

    AsyncCallback.DataCallback taskDataCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, byte[] data, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getTaskData((String) ctx);
                case OK:
                    List<String> workList = workerCache.getChildList();
                    if (CollectionUtils.isNotEmpty(workList)) {
                        int num = random.nextInt(workList.size());
                        String designatedWorker = workList.get(num);
                        String assignmentPath = ZooConstant.ASSIGN_PATH + "/" + designatedWorker + "/" + (String) ctx;
                        log.info("Assignment path: {}", assignmentPath);
                        createAssignment(assignmentPath, data);
                    }else{
                        log.error("No workers available");
                    }
            }
        }
    };

    private void createAssignment(String assignmentPath, byte[] data) {
        zk.create(assignmentPath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, assignTaskCallback, data);
    }

    //name就是create返回的创建路径
    AsyncCallback.StringCallback assignTaskCallback = new AsyncCallback.StringCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, String name) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    createAssignment(path, (byte[]) ctx);
                case OK:
                    log.info("createAssignment task : {}", name);
                    deleteTask(name.substring(name.lastIndexOf("/") + 1));

                case NODEEXISTS:
                    log.warn("Task already assigned");

                    break;
                default:
                    log.error("Error when trying to assign task.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    private void deleteTask(String task) {
        zk.delete(ZooConstant.TASK_PATH + "/" + task, -1, taskDeleteCallback, null);
    }


    private void getWorkers() {
        zk.getChildren(ZooConstant.WORKER_PATH, workersChangeWatcher, workersGetChildrenCallback, null);
    }



    Watcher workersChangeWatcher = (event) -> {
        if (event.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
            assert ZooConstant.WORKER_PATH.equals(event.getPath());
            getWorkers();
        }
    };

    AsyncCallback.ChildrenCallback workersGetChildrenCallback = (rc, path, ctx, children) -> {
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                getWorkers();
                break;
            case OK:
                log.info("Successfully got a list of workers: {}", children);
                reassignAndSet(children);
                break;
            default:
                log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
        }
    };

    private void reassignAndSet(List<String> children) {
        List<String> toProcess;
        if (workerCache == null) {
            workerCache = new ChildCache(children);
            toProcess = null;
        } else {
            toProcess = workerCache.removeChildDiff(children);
            log.info("removing and setting toProcess : {}", JSON.toJSONString(toProcess));
        }
        if (toProcess != null) {
            for (String worker : toProcess) {
                //处理失效的worker
                getAbsentWorkerTasks(worker);
            }
        }
    }

    private void getAbsentWorkerTasks(String worker) {
        zk.getChildren(ZooConstant.ASSIGN_PATH + "/" + worker, false, workerAssignmentCallback, null);
    }


    AsyncCallback.ChildrenCallback workerAssignmentCallback = new AsyncCallback.ChildrenCallback() {
        public void processResult(int rc, String path, Object ctx, List<String> children){
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getAbsentWorkerTasks(path);
                    break;
                case OK:
                    log.info("Succesfully got a list of assignments: "
                            + children.size()
                            + " tasks");
                    /*
                     * Reassign the tasks of the absent worker.
                     */

                    for(String task: children) {
                        getDataReassign(path + "/" + task, task);
                    }
                    break;
                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    //重新进行分配
    void getDataReassign(String path, String task) {
        //获取分配的task的信息
        zk.getData(path + "/" + task, false, getDataReassignCallback, task);
    }

    AsyncCallback.DataCallback getDataReassignCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, byte[] data, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getDataReassign(path, (String) ctx);
                case OK:
                    RecreateTaskCtx recreateTaskCtx = new RecreateTaskCtx(path, (String) ctx, data);
                    recreateTask(recreateTaskCtx);
                    break;
                default:
                    log.error("Something went wrong when getting data ", KeeperException.create(KeeperException.Code.get(rc)));
            }
        }
    };


    class RecreateTaskCtx {
        String path;
        String task;
        byte[] data;

        RecreateTaskCtx(String path, String task, byte[] data) {
            this.path = path;
            this.task = task;
            this.data = data;
        }
    }

    void recreateTask(RecreateTaskCtx ctx) {
        zk.create(ZkConstant.TASK_PATH + "/" + ctx.task,
                ctx.data,
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT,
                recreateTaskCallback,
                ctx);
    }

    AsyncCallback.StringCallback recreateTaskCallback = new AsyncCallback.StringCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, String name) {
            RecreateTaskCtx taskCtx = (RecreateTaskCtx) ctx;
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    recreateTask(taskCtx);
                    break;
                case OK:
                    deleteAssignment(taskCtx.path);
                    break;
                default:
                    log.error("Something wwnt wrong when recreating task", KeeperException.create(KeeperException.Code.get(rc)));
            }
        }
    };

    void deleteAssignment(String path){
        zk.delete(path, -1, taskDeleteCallback, null);
    }

    AsyncCallback.VoidCallback taskDeleteCallback = new AsyncCallback.VoidCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx) {
            switch(KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    deleteAssignment(path);
                    break;
                case OK:
                    log.info("Task correctly deleted: {}", path);
                    break;
                default:
                    log.error("Failed to delete task data" + KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };
    private void masterExists() {
        zk.exists(ZooConstant.MASTER_PATH, masterExistWatcher, masterExistCallback, null);
    }

    private Watcher masterExistWatcher = (event) -> {
        if (event.getType() == Watcher.Event.EventType.NodeDeleted) {
            assert ZooConstant.MASTER_PATH.equals(event.getPath());
            runForMaster();
        }
    };

    private AsyncCallback.StatCallback masterExistCallback = (rc, path, ctx, stat) -> {
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                masterExists();
                break;
            case OK:
                try {
                    byte[] data = zk.getData(ZooConstant.MASTER_PATH, false, null);
                    this.leaderId = new String(data);
                    log.info("my name : {} and the lead is {} ", name, leaderId);
                } catch (Exception e) {
                    log.error("", e);
                }
                break;
            case NONODE:
                runForMaster();
                log.info("It sounds like the previous master is gone, " +
                        "so let's run for master again.");
            default:
                checkMaster();
                break;
        }
    };



    private void checkMaster() {
        zk.getData(ZooConstant.MASTER_PATH, false, masterCheckCallback, null);
    }

    private AsyncCallback.DataCallback masterCheckCallback = (rc, path, ctx, data, stat) -> {
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                checkMaster();
                return;
            case NONODE:
                runForMaster();
                return;
            case OK:
                if (name.equals(new String(data))) {
                    log.info("I am master , name = {} ", name);
                    this.leaderId = name;
                    status = Status.ELECTED;
                    takeLeaderShip();
                } else {
                    status = Status.NOTELECTED;
                    leaderId = new String(data);
                    log.info("master is {}", leaderId);
                    masterExists();
                }
                break;
            default:
                log.error("something went wrong", KeeperException.create(KeeperException.Code.get(rc), path));
        }
    };



    public static void main(String[] args) {
        Master master = new Master("localhost:2181");
        try {
            master.startZK();
            master.init();
            master.runForMaster();
            log.info("master name : {}", master.name());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public String leaderId() {
        return leaderId;
    }
}